import ROOT as rt
from tqdm import tqdm

def LineShapeBifurGaus(x, par):
    arg = x[0] - par[1]
    if arg < 0:
        return  par[0] * rt.TMath.Gaus(x[0], par[1], par[2])
    else:
        return  par[0] * rt.TMath.Gaus(x[0], par[1], par[3])

def LineShapeNovosibirks(x, par):
    peak = par[1]
    width = par[2]
    tail = par[3]

    if (rt.TMath.Abs(tail) < 1.e-7):
        return rt.TMath.Exp( -0.5 * rt.TMath.Power( ( (x[0] - peak) / width ), 2 ))

    arg = 1.0 - ( x[0] - peak ) * tail / width

    if arg < 1.e-7:
    #Argument of logarithm negative. Real continuation -> function equals zero
        return 0.0
    log = rt.TMath.Log(arg)
    xi = 2.3548200450309494 # 2 Sqrt( Ln(4) )

    width_zero = ( 2.0 / xi ) * rt.TMath.ASinH( tail * xi * 0.5 )
    width_zero2 = width_zero * width_zero
    exponent = ( -0.5 / (width_zero2) * log * log ) - ( width_zero2 * 0.5 )

    return par[0] * rt.TMath.Exp(exponent)

def LineShapeBukin(x, par):
    import math
    consts = 2*math.sqrt(2*math.log(2.0))
    Xp = par[1]
    sigp = par[2]
    xi = par[3]
    rho1 = par[4]
    rho2 = par[5]

    hp=sigp*consts
    r3=math.log(2.)
    r4=math.sqrt(xi*xi+1)
    r1=xi/r4

    if math.fabs(xi) > math.exp(-6.) :
        r5=xi/math.log(r4+xi)
    else:
        r5=1

    x1 = Xp + (hp / 2) * (r1-1)
    x2 = Xp + (hp / 2) * (r1+1)

  #--- Left Side
    if x[0] < x1:
        r2=rho1*(x[0]-x1)*(x[0]-x1)/(Xp-x1)/(Xp-x1)-r3 + 4 * r3 * (x[0]-x1)/hp * r5 * r4/(r4-xi)/(r4-xi)
    
  #--- Center
    elif x[0] < x2:
        if math.fabs(xi) > math.exp(-6.) :
            r2=math.log(1 + 4 * xi * r4 * (x[0]-Xp)/hp)/math.log(1+2*xi*(xi-r4))
            r2=-r3*r2*r2
        else:
            r2=-4*r3*(x[0]-Xp)*(x[0]-Xp)/hp/hp

  #--- Right Side
    else:
        r2=rho2*(x[0]-x2)*(x[0]-x2)/(Xp-x2)/(Xp-x2)-r3 - 4 * r3 * (x[0]-x2)/hp * r5 * r4/(r4+xi)/(r4+xi)

    if math.fabs(r2) > 100:
        fit_result = 0
    else:
    #---- Normalize the result
        fit_result = math.exp(r2)
  
    return par[0] * fit_result

def DrawLineShape(args):
    from array import array

    rt.gROOT.ProcessLine(".L lhcbStyle.C")
    inFile = rt.TFile(args.InputFile)
    inTree = inFile.Get("SPE")

    maxChi2 = 100000
    cut = f"IsValid==1&&FitChi2<{maxChi2}"

    print(f"Total Entries: {inTree.GetEntries()}")
    print(f"Accept: {inTree.GetEntries(cut)}")

    thTau = rt.TH1F("thTau", "", 100, 1., 5.)
    inTree.Project("thTau", "Tau", cut)
    thTau.SetLineColor(rt.kBlack)
    fitTau = rt.TF1("fitTau", LineShapeBukin, 1., 5., 6)
    fitTau.SetParameters(800., 2.8, 0.3, 0.3, -0.3, 0.06)
    fitTau.SetLineColor(rt.kRed)
    fitTau.SetLineWidth(1)
    resultTau = thTau.Fit(fitTau, "S", "", 1., 5.)

    thSmear = rt.TH1F("thSmear", "", 100, 0.5, 2.)
    inTree.Project("thSmear", "Smear", cut)
    thSmear.SetLineColor(rt.kBlack)
    fitSmear = rt.TF1("fitSmear", LineShapeBukin, 0., 5., 6)
    fitSmear.SetParameters(1400., 1.0, 0.1, 0.1, -0.1, 0.06)
    fitSmear.SetLineColor(rt.kRed)
    fitSmear.SetLineWidth(1)
    resultSmear = thSmear.Fit(fitSmear, "S", "", 0., 5.)

    thAmp = rt.TH1F("thAmp", "", 100, 0., 7000.)
    inTree.Project("thAmp", "Amp", cut)
    thAmp.SetLineColor(rt.kBlack)
    fitAmp = rt.TF1("fitAmp", "[0]*TMath::Gaus(x,[1],[2])", 2000., 5000.)
    fitAmp.SetParameters(12000., 2900., 974)
    fitAmp.SetParLimits(1, 2500, 3200)
    fitAmp.SetParLimits(2, 500, 1000)
    fitAmp.SetParLimits(3, 0., 1.)
    fitAmp.SetLineColor(rt.kRed)
    fitAmp.SetLineWidth(1)
    resultAmp = thAmp.Fit(fitAmp, "S", "", 2000., 3500.)

    thQ_Amp = rt.TH2F("thQ_Amp", "", 100, 0, 5500, 100, 0, 5500)
    inTree.Project("thQ_Amp", "sigq_ch0:Amp", cut)
    fitLine = rt.TF1("fitLine", "x*[0]", 0, 6000)
    fitLine.SetLineColor(rt.kRed)
    fitLine.SetLineWidth(1)
    resultQ_Amp = thQ_Amp.Fit(fitLine, "S", "", 0, 5500)

    can = rt.TCanvas("can", "", 1600, 1200)
    can.Divide(2,2)

    tex = rt.TLatex()
    tex.SetTextAlign(12)
    tex.SetTextSize(0.055)
    tex.SetNDC()

    can.cd(1)
    thTau.GetXaxis().SetTitle("\\tau")
    thTau.Draw("E")
    fitTau.Draw("SAME")
    tex.DrawLatex(0.6, 0.8, "\\mu = {:.3f} \\pm {:.3f}".format(resultTau.GetParams()[1], resultTau.GetErrors()[1]))
    tex.DrawLatex(0.6, 0.7, "\\sigma_{{L}} = {:.3f} \\pm {:.3f}".format(resultTau.GetParams()[2], resultTau.GetErrors()[2]))
    tex.DrawLatex(0.6, 0.6, "\\sigma_{{R}} = {:.3f} \\pm {:.3f}".format(resultTau.GetParams()[3], resultTau.GetErrors()[3]))

    can.cd(2)
    thSmear.GetXaxis().SetTitle("\\sigma")
    thSmear.Draw("E")
    fitSmear.Draw("SAME")
    tex.DrawLatex(0.6, 0.8, "\\mu = {:.3f} \\pm {:.3f}".format(resultSmear.GetParams()[1], resultSmear.GetErrors()[1]))
    tex.DrawLatex(0.6, 0.7, "\\sigma_{{L}} = {:.3f} \\pm {:.3f}".format(resultSmear.GetParams()[2], resultSmear.GetErrors()[2]))
    tex.DrawLatex(0.6, 0.6, "\\sigma_{{R}} = {:.3f} \\pm {:.3f}".format(resultSmear.GetParams()[3], resultSmear.GetErrors()[3]))

    can.cd(3)
    thAmp.GetXaxis().SetTitle("Q")
    thAmp.Draw("E")
    fitAmp.Draw("SAME")
    tex.DrawLatex(0.6, 0.8, "\\mu = {:.3f} \\pm {:.3f}".format(resultAmp.GetParams()[1], resultAmp.GetErrors()[1]))
    tex.DrawLatex(0.6, 0.7, "\\sigma_{{L}} = {:.3f} \\pm {:.3f}".format(resultAmp.GetParams()[2], resultAmp.GetErrors()[2]))
    tex.DrawLatex(0.6, 0.6, "\\sigma_{{R}} = {:.3f} \\pm {:.3f}".format(resultAmp.GetParams()[3], resultAmp.GetErrors()[3]))

    can.cd(4)
    thQ_Amp.GetXaxis().SetTitle("Amp")
    thQ_Amp.GetYaxis().SetTitle("Q")
    thQ_Amp.Draw("colz")
    fitLine.Draw("SAME")
    tex.DrawLatex(0.5, 0.3, "k = {:.3f} \\pm {:.3f}".format(resultQ_Amp.GetParams()[0], resultQ_Amp.GetErrors()[0]))
    can.SaveAs(f"{args.SavePath}/FitResult.pdf")

def main():
    import argparse
    import time
    import os

    parser = argparse.ArgumentParser(description="Significance")
    parser.add_argument(
        "-i", "--InputFile", default="SPE_waveform_fitted.root", dest="InputFile"
    )
    parser.add_argument(
        "-s", "--SavePath", default="Fits", dest="SavePath"
    )
    args = parser.parse_args()

    if not os.path.exists(f"{args.SavePath}"):
        os.makedirs(f"{args.SavePath}")

    DrawLineShape(args)


if __name__ == "__main__":
    main()